library(data.table)
library(mefa)
library(doParallel)
library(emayili)
library(magrittr)


source('helper_functions.r')


runDemandEstimation <- function(pars,filename) { 
  
  # For logging
  if(!pars$DEV) {
    sink(file = paste0('estimation_logs/',filename,'.log'),append = T)
  }
  
  if(F) {
    pars <- params
    pars$EST_YEARS      <- 2010:2017
  }
  
  # Load data and create raw draws
  d.in <- fread('../data/final/estimation_data_oct_29.csv')
  if(pars$DEV) {
    all.markets <- unique(d.in$MARKET_ID)
    markets.data <- d.in[MARKET_ID %in% all.markets[1:100]]
  } else {
    markets.data <- d.in[YEAR %in% pars$EST_YEARS]
  }
  
  # Create instrumented interest rate
  rate.model <- lm(RATE ~ JUMBO * (D60.L1 ) + as.factor(paste0(YEAR,CBSA))   , data = markets.data,weights=markets.data$MARKET_SIZE)
  
  markets.data[,RATE_INST := predict(rate.model,markets.data)]
  
  # Set initial guess and create parameters
  x0 = pars$initial_guess
  new_mx <- param_vec_to_matrix(x0)
  pars$A <- new_mx$A
  pars$B <- new_mx$B
  pars$S <- new_mx$S
  raw.draws <- createRawDraws(pars)
  delta.init <<- NULL
  
  # Pre-loop
  # Calculate weights (if needed)
  # Initialize delta and drop some unsolveable markets.
  
  temp.pars <- pars
  temp.pars$BLP_TOL = 1e-10
  temp.pars$BLP_ATTEMPTS = 1000
  contraction <- solveAllMarkets(markets.data,raw.draws,temp.pars,delta.init,cores = temp.pars$CORES)
  print(paste0('Dropping ',sum(contraction$converged == 0)/4,'/',nrow(contraction)/4,' markets.'))
  contraction <- contraction[converged == 1]
  markets.data <- markets.data[MARKET_ID %in% contraction$MARKET_ID]
  delta.init <<- contraction[,c('MARKET_ID','j','delta')]
  
  if(pars$moment_weights == 'inverse_mean') {
    market.data.merged <- merge(markets.data,contraction,by=c('MARKET_ID','j'))
    if(length(pars$EST_YEARS) > 1) {
      linear <- lm(delta ~ RATE_INST + JUMBO + as.factor(YEAR) * PURPOSE * LENDER,weights = market.data.merged$converged,data = market.data.merged)
    } else {
      linear <- lm(delta ~ RATE_INST + JUMBO + PURPOSE * LENDER,weights = market.data.merged$converged,data = market.data.merged)
    }
    moments <- calculateMoments(market.data.merged,linear)
    moment_weights <- data.table(moment = moments$moment,weight = 1 / moments$value^2)
    moment_weights[moment == 'convergence.pct']$weight <- 1
  } else {
    moment_weights <- pars$moment_weights
  }
  
  
  # Calculate GMM objective funcion
  gmm_objective <- function(x,diag = F) {
    
    
    # 1.  Deal parameters
    new_mx <- param_vec_to_matrix(x)
    pars$A <- new_mx$A
    pars$B <- new_mx$B
    pars$S <- new_mx$S
    
    if(pars$diag.show_parameters) {
      print(new_mx)
    }
    
    # 2. Recover deltas
    contraction <- solveAllMarkets(markets.data,raw.draws,pars,delta.init,cores = pars$CORES)
    
    # 3. Update the initial guesses for the good ones.
    next.init <- contraction
    next.init[converged == 0,delta := mean(next.init[converged == 1]$delta,na.rm=T)]
    delta.init <<- next.init[,c('MARKET_ID','j','delta')]
    
    # 4. Merge the deltas to the main marketdata
    market.data.merged <- merge(markets.data,contraction,by=c('MARKET_ID','j'))
    
    # 5. Regress out the remaining parameters
    if(length(pars$EST_YEARS) > 1) {
      linear <- lm(delta ~ RATE_INST + JUMBO + as.factor(YEAR) * PURPOSE * LENDER,weights = market.data.merged$converged,data = market.data.merged)
    } else {
      linear <- lm(delta ~ RATE_INST + JUMBO + PURPOSE * LENDER,weights = market.data.merged$converged,data = market.data.merged)
    }
    
    # 6. Calculate moments, apply wieghts, and calculate objective function
    moments <- calculateMoments(market.data.merged,linear)
    moments[,value := model - data]
    moments <- merge(moments,moment_weights,by='moment')
    val     <- sum(moments$value^2 * moments$weight)
    
    # 7. Report diagnostics if we want
   if(pars$diag.show_linear) {
      print(linear)
    }
    if(pars$diag.show_moments) {
      moments[,contribution := value^2 * weight]
      print(moments)  
    }
    if(pars$diag.show_convergence) {
      print(paste0(sum(contraction$converged)/4,'/',nrow(contraction)/4,' markets converged. Mean iterations given convergence = ',mean(contraction[converged == 1]$attempts)))
    }
    
    
    # 8. Report actual objective function
    print(val)
    
    if(!diag) {
      return(val)
    } else {
      return(list(x = x, market.data.merged = market.data.merged, linear.model = linear, moments = moments, val = val))
    }  
    
    
    
  }
  
  # Run the optimization
  scale = c(1,1,.1,.1,.1,.1,.1,.1,.1,.1,.1,.1,.1,.1)
  estimation <- optim(par = x0, fn = gmm_objective,method = 'L-BFGS-B',lower = c(1,11,-1,-3,-4,0,-1,-3,-4,0,.01,.01,.01,.01),control = list(parscale = scale,lmm = 10))
  

  # Solve the model and return diagnostics to save
  solution <- gmm_objective(estimation$par,diag = T)
  estimation_data = list(pars = pars, market.data = markets.data, raw.draws = raw.draws, estimation = estimation,solution = solution,moment_sensitivity = JJ)
  save(list = c('estimation_data'),file = paste0('estimation_results/',filename,'.rsave'))
}


## Execute it.

if(F) {
  # Load config parameters
  config = commandArgs(trailingOnly = T)
  #config = 'estimation_24'
  #source(paste0('estimation_configs/',config,'.r'))
  
  # Save all the sourced functions.
  save(list = lsf.str(),file = paste0('estimation_functions/',pars$filename,'.rsave'))
  
  ## Run the estimation
  runDemandEstimation(pars,pars$filename)
}



